#!/usr/bin/env python -Wignore
# encoding: utf-8
#
# Copyright (c) 2012 Philip Jameson and Daniel Selans
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without 
# modification, are permitted provided that the following conditions are met:
#
#  - Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation 
#    and/or other materials provided with the distribution.
#  - Neither the name of the <ORGANIZATION> nor the names of its contributors 
#    may be used to endorse or promote products derived from this software 
#    without specific prior written permission.
#
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF 
# THE POSSIBILITY OF SUCH DAMAGE.
#

"""
xmlserver.py -- a threaded/async and ssl'd xmlrpc server based on the built in XMLRPC Server 
"""

import SocketServer
import BaseHTTPServer
import SimpleHTTPServer
import SimpleXMLRPCServer

import logging
import socket
import select
import base64

from OpenSSL import SSL


class XMLRPCRegisters:
	""" XMLRPC callback functions """
	def ping(self):
		return 'pong'
	pass

class XMLServer:
	"""
	Main XML server class 

	This is the main server class that should be instansiated. It sets up the XMLRPC server and
	all of the required classes for SSL

	Attributes:
		handler_class - The class that represents the methods available for use via XMLRPC.
				if not specified, it is XMLRPCRegisters. This cannot be changed after
				'start()' is called on your XMLServer instance
		logger        -	The logger object to use. If not given, a Null logger is setup. This object
			 	should implment the logging functions found in objects from 'logger.getLogger()'
		allowed_ips   - A list of strings that contain IPs that are allowed to connect to the server. If
				set to None, all IPs are allowed
		require_auth  - Whether or not we require HTTP authorization headers or not. Basic is the only method supported
	"""
	handler_class = None
	logger = None
	allowed_ips = ['127.0.0.1']
	require_auth = False
	_user_info = {}
	"""
	Constructor for the XMLServer class. Setup the entire server here

	Arguments:
		host 		- String with the address to listen on. Normally 0.0.0.0
		port		- The port to listen on. This should be an integer
		logger		- A logger object to use for error messages. Default is None
		handler_class	- The class to use for XMLRPC functions. Explained above
		allowed_ips	- A list of strings of IPs to allow connections from. None if you want all
				  IPs to access the server 
		cert_file	- A filename where the server's certificate is located
		key_file	- A filename where the server's key is located
	"""
	def __init__(self, host, port, logger = None, handler_class = XMLRPCRegisters, allowed_ips = ['127.0.0.1'], cert_file = 'cert.pem', key_file = 'key.pem'):

		self.server_address = host
		self.server_port = port
		self.cert_file = cert_file
		self.key_file = key_file

		if logger == None:
			logger = logging.getLogger('XMLServer')
			logger.addHandler(logging.NullHandler)
		self.logger = logger
		self.handler_class = handler_class
		self.allowed_ips = allowed_ips

	"""
	Starts the server in a blocking fashion. Requests, though, are dispatched asynchronously
	"""
	def start(self):
		server_address = (self.server_address, self.server_port)
		server = AsyncXMLRPCServer(server_address, XMLRPCHandler, self, self.cert_file, self.key_file)
		server.register_instance(self.handler_class())
		sa = server.socket.getsockname()
		try:
			server.serve_forever()
		except socket.error,e:
			pass
		except select.error,e:
			pass
		finally:
			server.server_close()

	"""
	Add a username/passowrd combination to the allowed list
	"""
	def add_user(self, username, password):
		self._user_info[username] = password

	"""
	Remove a username/password combination from the allowed list
	"""
	def remove_user(self, username):
		if username in self._user_info:
			del self._user_info[username]

	"""
	Check whether a username and password are correct
	
	Arguments:
		username - The username
		password - The password
	
	Returns:
		True if a valid user, false if not
	"""
	def check_auth(self, username, password):
		if username in self._user_info:
			if self._user_info[username] == password:
				return True
		return False

class SecureXMLRPC(BaseHTTPServer.HTTPServer, SimpleXMLRPCServer.SimpleXMLRPCDispatcher):
	""" Based on Laszlo Nagy's Secure XMLRPC implementation """
	def __init__(self, server_address, HandlerClass, logRequests=False, cert_file = 'cert.pem', key_file = 'key.pem'):
		self.logRequests = logRequests
		SimpleXMLRPCServer.SimpleXMLRPCDispatcher.__init__(self)
		SocketServer.BaseServer.__init__(self, server_address, HandlerClass)
		ctx = SSL.Context(SSL.SSLv23_METHOD)
		ctx.use_privatekey_file(key_file)
		ctx.use_certificate_file(cert_file)
		self.socket = SSL.Connection(ctx, socket.socket(self.address_family, self.socket_type))
		self.server_bind()
		self.server_activate()

class XMLRPCHandler(SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
        def setup(self):
        	self.connection = self.request
        	self.rfile = socket._fileobject(self.request, "rb", self.rbufsize)
        	self.wfile = socket._fileobject(self.request, "wb", self.wbufsize)

	def do_GET(self):
		try:
			self.send_error(414)
			self.end_headers()
		except:
			self.server.xml_server.logger.critical('Could not send HTTP status 414 after GET request to %s' % self.connection.getpeername()[0])
		self.wfile.flush()
		self.connection.shutdown()
		self.connection.close()

        def do_POST(self):
		client_ip = self.connection.getpeername()[0]
        	if self.server.xml_server.allowed_ips != None and client_ip not in self.server.xml_server.allowed_ips:
			self.server.xml_server.logger.critical("Unauthorized connection from IP %s" % client_ip)
        		self.wfile.flush()
        		self.connection.shutdown()
			self.connection.close()
        		return

		if self.server.xml_server.require_auth == True:
			error = None
			if 'authorization' not in self.headers:
				self.server.xml_server.logger.warning('Got unauthenticated request from %s when authentication is required' % client_ip)
				error = 401
			else:
				parts = self.headers['authorization'].split()
				if len(parts) != 2 or parts[0].lower() != 'basic':
					error = 401
				else:
					try:
						decoded = base64.b64decode(parts[1])
						parts = decoded.split(':', 1)
						if len(parts) != 2:
							self.server.xml_server.logger.warning('Got invalid authorization packet from %s (no ":")' % client_ip)
							error = 400
						else:
							username = parts[0]
							password = parts[1]
					except TypeError, e:
						self.server.xml_server.logger.warning('Got invalid authorization packet from %s' % client_ip)
						error = 400
			
					if error == None:
						if not self.server.xml_server.check_auth(username, password):
							error = 401

			if error != None:
				self.send_response(error)
				if error == 401:
					self.send_header('WWW-Authenticate', 'Basic realm="XMLRPC Service"')
				self.end_headers()
				self.wfile.flush()
				self.connection.shutdown()
				return

        	try:
        		data = self.rfile.read(int(self.headers["content-length"]))
        		response = self.server._marshaled_dispatch(data, getattr(self, '_dispatch', None))
        		# Fairly heavy
			self.server.xml_server.logger.debug("New request from %s: %s" % (client_ip, data.replace('\n', '')))
        	except Exception, e:
			self.server.xml_server.logger.critical("Exception: %s" % e)
        		self.send_response(500)
        		self.end_headers()
        	else:
        		try:
        			self.send_response(200)
        			self.send_header("Content-type", "text/xml")
        			self.send_header("Content-length", str(len(response)))
        			self.end_headers()
        			self.wfile.write(response)
        		except:
				self.server.xml_server.logger.critical("Unable to send response to client %s" % client_ip)
        		
        		self.wfile.flush()
        		self.connection.shutdown()
			self.connection.close()


class AsyncXMLRPCServer(SocketServer.ThreadingMixIn, SecureXMLRPC):
	""" Extending SocketServer.ThreadingMixIn to add async capabilities to SecureXMLRPC"""
	#Instance of XMLServer that we can get/set properties from XMLRPCHandler's instance
	xml_server = None

	def __init__(self, server_address, HandlerClass, xml_server_instance, cert_file = 'cert.pem', key_file = 'key.pem'):
		self.xml_server = xml_server_instance
		SecureXMLRPC.__init__(self, server_address, HandlerClass, logRequests=False, cert_file = cert_file, key_file = key_file)
